import torch
from torch import nn
import json
import numpy as np

class StockNetS(nn.Module):
    def __init__(self):
        super().__init__()
        self.r1 = nn.RNN(10, 20, 1)
        self.r2 = nn.RNN(20, 20, 2)
        self.r3 = nn.RNN(20, 10, 1)
        self.l1 = nn.Linear(10, 10)

    def forward(self, x):
        x, _ = self.r1(x)
        x, _ = self.r2(x)
        x, _ = self.r3(x)
        x = self.l1(x)
        return x

class StockNetL(nn.Module):
    def __init__(self):
        super().__init__()
        self.r1 = nn.LSTM(100, 110, 1)
        self.r2 = nn.LSTM(110, 110, 1)
        self.r3 = nn.LSTM(110, 10, 1)
        self.l1 = nn.Linear(10, 10)

    def forward(self, x):
        x, _ = self.r1(x)
        x, _ = self.r2(x)
        x, _ = self.r3(x)
        x = self.l1(x)
        return x

def runRnn():
    model = StockNetS()
    a = model(torch.rand(1, 1, 10)).detach().numpy()[0][0:7]
    b = {"data":a.tolist()[0]}
    with open("./infer.json", "w") as f:
        json.dump(b, f)

if __name__ == "__main__":
    runRnn()